# chatgpt
import os
import argparse
from locust import HttpUser, task, between, events
from config.all_paths import get_root_path
from config.global_variable import ENV_DICT
from utils.file_handling.excel_handling import ExcelUtil
from utils.interface_auth.auth_tool import Authentication
from utils.performance_tool.config_loader import load_config
from utils.performance_tool.load_shapes import SHAPE_MAP
from utils.request_tool.send_request import RequestSender
from utils.time_tool.time_control import parse_duration

# ====================== 命令行参数处理 ======================
parser = argparse.ArgumentParser()
parser.add_argument("--env", type=str, help="目标环境，如 test_wt", required=True)
parser.add_argument("--role", type=str, help="用户角色，如 doctor", required=True)
parser.add_argument("--profile_name", type=str, help="负载配置名称，如 stress", required=True)
args, unknown = parser.parse_known_args()

# 将参数传入环境变量，供 locust 初始化时使用
os.environ["LOCUST_ENV"] = args.env
os.environ["LOCUST_ROLE"] = args.role
os.environ["LOCUST_PROFILE"] = args.profile_name

# ====================== 读取测试数据 ======================
excel_name = 'data/case_excel/wutong/uat/医生账号相关接口.xlsx'
test_data = ExcelUtil(os.path.join(get_root_path(), excel_name)).read_excel()

# ====================== 加载配置 ======================
env = os.environ.get("LOCUST_ENV", "test_wt")
role = os.environ.get("LOCUST_ROLE", "doctor")
profile_name = os.environ.get("LOCUST_PROFILE", "baseline")

config = load_config(env, role)
profile_config = config['profiles'].get(profile_name, {})

# ====================== 测试生命周期钩子 ======================
@events.test_start.add_listener
def init_test(environment, **kwargs):
    """测试初始化"""
    print(f"\n=== 启动 {profile_name} 测试 ===")
    print(f"目标环境: {env}")
    print(f"基准URL: {config['env'].get('url')}")


@events.test_stop.add_listener
def finalize_test(environment, **kwargs):
    """测试收尾"""
    print("\n=== 测试结果摘要 ===")
    print(f"总请求数: {environment.stats.total.num_requests}")
    print(f"失败率: {environment.stats.total.fail_ratio:.2%}")

# ====================== 用户类定义 ======================
class BaseUser(HttpUser):
    wait_time = between(0.5, 2)
    abstract = True

    def on_start(self):
        """根据环境选择不同认证方式"""
        auth = Authentication(
            ENV_DICT.get('username'),
            ENV_DICT.get('password'),
            ENV_DICT.get('url'),
        )

        auth_method_map = {
            "test_wt": "auth_wutong",
            "test_th": "auth_th",
            "test_inter": "auth_inter"
        }

        env_key = os.environ.get("LOCUST_ENV", "test_wt")
        auth_method_name = auth_method_map.get(env_key, "auth_wutong")

        if hasattr(auth, auth_method_name):
            self.s = getattr(auth, auth_method_name)()
            print(f"✅ 使用 {auth_method_name} 进行认证")
        else:
            raise ValueError(f"❌ 未找到对应的认证方法: {auth_method_name}")


class BaselineUser(BaseUser):
    """恒定负载用户"""
    @task
    def health_check(self):
        RequestSender().send_request(self.s, test_data[0])


class StressUser(BaseUser):
    """峰值压力测试用户"""
    @task(5)
    def create_data(self):
        RequestSender().send_request(self.s, test_data[0])


class ConcurrencyUser(BaseUser):
    """瞬时并发用户"""
    @task
    def update_resource(self):
        resource_id = profile_config.get("concurrency_key")
        self.client.put(f"/resources/{resource_id}")


class WaveUser(BaseUser):
    """波浪型负载用户"""
    @task
    def wave_task(self):
        self.client.get("/api/wave-test")


# ====================== 动态测试类选择 ======================
class DynamicTestClass(eval(profile_config.get("user_class", "BaselineUser"))):
    pass

# ====================== 自定义负载形状 ======================
class CustomLoadShape(SHAPE_MAP[profile_config.get("load_shape", "constant")]):
    shape_users = profile_config.get("users", 1)
    spawn_rate = profile_config.get("spawn_rate", 1)
    shape_duration = parse_duration(profile_config.get("duration", "1m"))
    shape_stages = profile_config.get("stages", [])

# ====================== 供 Locust 使用的任务类和形状类 ======================
class TestScenario(DynamicTestClass):
    pass

def get_shape_class():
    return CustomLoadShape
